#include <libwebsockets.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <time.h>
#include <stdbool.h>
#include <pthread.h>

/** =========================================================================================================== */
/** Snowflake ID generation algorithm */
#include <stdint.h>

/** Define the structure of Snowflake ID */
#define TIMESTAMP_BITS 41
#define WORKER_ID_BITS 10
#define SEQUENCE_BITS 12
#define MAX_TIMESTAMP ((1L << TIMESTAMP_BITS) - 1)
#define MAX_WORKER_ID ((1L << WORKER_ID_BITS) - 1)
#define MAX_SEQUENCE ((1L << SEQUENCE_BITS) - 1)
#define TIMESTAMP_SHIFT (SEQUENCE_BITS + WORKER_ID_BITS)
#define WORKER_ID_SHIFT (SEQUENCE_BITS)
#define MSG_ID_STR_SIZE 20
/** Initial timestamp (milliseconds of 00:00:00 on December 1, 2024) */
#define EPOCH 1706627200000L

/** global variable for Snowflake ID */
static uint64_t last_timestamp = 1L;
static uint64_t sequence = 0L;
static uint64_t worker_id = 1L;

/** Get the current time in milliseconds */
uint64_t current_millis() {
  struct timespec ts;
  clock_gettime(CLOCK_REALTIME, &ts);
  return (uint64_t)ts.tv_sec * 1000L + (uint64_t)ts.tv_nsec / 1000000L;
}

/** wait for next micoseconds */
uint64_t wait_next_millis(uint64_t last_timestamp) {
  uint64_t timestamp = current_millis();
  while (timestamp <= last_timestamp) {
    timestamp = current_millis();
  }
  return timestamp;
}

/** generate Snowflake ID */
uint64_t snowflake_id() {
  uint64_t timestamp = current_millis();
  if (timestamp < last_timestamp) {
    fprintf(stderr, "Clock moved backwards. Refusing to generate id.\n");
    return -1;
  }
  if (timestamp == last_timestamp) {
    sequence = (sequence + 1) & MAX_SEQUENCE;
    if (sequence == 0) {
      /** Serial number overflow, waiting for the next millisecond */
      timestamp = wait_next_millis(last_timestamp);
    }
  } else {
    sequence = 0L;
  }
  last_timestamp = timestamp;
  uint64_t id = ((timestamp - EPOCH) << TIMESTAMP_SHIFT) |
    (worker_id << WORKER_ID_SHIFT) |
    sequence;
  return id;
}
/**  =========================================================================================================== */

#define MAX_A_CLIENTS 10
#define CONFIG_FILE "./msg-router.conf"
#define LOG_FILE "./log/msg-router.log"
#define LWS_RX_BUF_SIZE 4096
#define LWS_TIMEOUT_MS  50
#define LOG_MSG_BUF_SIZE  256

typedef struct {
  struct lws *wsi;
  int client_type;
  uint64_t message_id;
  /** 0: no message; 1: queue; 2: processing */
  int msg_status;
  char * msg_cache;
  size_t len;
} client_data_t;

/** global variable for Websocket Server */
static client_data_t *a_clients[MAX_A_CLIENTS] = {0};
static client_data_t *b_client = NULL;
volatile sig_atomic_t terminate_server = 0;
static bool waitting_for_gpt = false;
pthread_mutex_t clients_mutex = PTHREAD_MUTEX_INITIALIZER;

void log_to_file(const char *message) {
  FILE *file = fopen(LOG_FILE, "a");
  if (file) {
    char timestamp[64];
    memset(timestamp, 0, sizeof(timestamp));
    time_t now = time(NULL);
    struct tm *tm_info = localtime(&now);
    strftime(timestamp, sizeof(timestamp), "%Y-%m-%d %H:%M:%S", tm_info);
    fprintf(file, "[%s] %s\n", timestamp, message);
    fclose(file);
  }
}

static bool send_msg(struct lws *wsi, void *msg, size_t len){
  unsigned char *buffer = malloc(LWS_PRE+len);
  if(NULL == buffer){
    log_to_file("malloc failed in send_msg");
    return false;
  }
  memset(buffer, 0, LWS_PRE+len);
  memcpy(buffer + LWS_PRE, msg, len);
  int n = lws_write(wsi, buffer + LWS_PRE, len, LWS_WRITE_TEXT);
  if (n < 0) {
    log_to_file("Failed to send message to client");
    return false;
  }
  free(buffer);
  return true;
}

uint64_t peek_msg_id(char * msg){
  unsigned char tmp[21];
  strncpy(tmp, msg, 20);
  tmp[20] = '\0';
  return strtoull(tmp, NULL, 10);
}

bool disconnect(struct lws *wsi, unsigned char *log_msg, unsigned char *msg){
  size_t len = strlen(msg);
  unsigned char * buffer = malloc(len+LWS_PRE);
  if(NULL == buffer){
    log_to_file("malloc failed in disconnect");
    return false;
  }
  memset(buffer, 0, len+LWS_PRE);
  memcpy(buffer+LWS_PRE, msg, len);
  lws_close_reason(wsi, LWS_CLOSE_STATUS_NORMAL, buffer+LWS_PRE, len);
  lws_set_timeout(wsi, LWS_TO_KILL_ASYNC, 0);
  free(buffer);
  return true;
}

static int wss_a_callback(struct lws *wsi, enum lws_callback_reasons reason, void *user, void *in, size_t len) {
  client_data_t *client_data = (client_data_t *)user;
  char log_message[LOG_MSG_BUF_SIZE];
  memset(log_message, 0, sizeof(log_message));
  switch (reason) {
    case LWS_CALLBACK_ESTABLISHED:
      pthread_mutex_lock(&clients_mutex);
      bool is_full = true;
      for (int i = 0; i < MAX_A_CLIENTS; ++i) {
        if (a_clients[i] == NULL) {
          a_clients[i] = client_data;
          a_clients[i]->wsi = wsi;
          a_clients[i]->client_type = 0;
          sprintf(log_message, "A client connected at slot %d", i);
          log_to_file(log_message);
          is_full = false;
          break;
        }
      }
      if(is_full){
        /** disconnect */
        sprintf(log_message, "A client rejected (there is no slot for the client)");
        unsigned char * msg = "Only one 10 client allowed";
        disconnect(wsi, log_message, msg);
        pthread_mutex_unlock(&clients_mutex);
        return -1;
      }
      pthread_mutex_unlock(&clients_mutex);
      break;
    case LWS_CALLBACK_RECEIVE:
      pthread_mutex_lock(&clients_mutex);
      if(b_client && b_client->wsi) {
        if(client_data->msg_status == 1 || client_data->msg_status == 2){
          /** tell the client that it has message need replay */
          unsigned char * msg = "You alse have a message in processing!";
          send_msg(wsi, msg, strlen(msg));
        }else if(waitting_for_gpt){ /** there is a message in processing cache the message */
          client_data->message_id = snowflake_id();
          client_data->msg_cache = malloc(len);
          memcpy(client_data->msg_cache, in, len);
          client_data->msg_status = 1;
          client_data->len = len;

          /** tell the client that there is a message in processing */
          unsigned char * msg = "There is a message in processing!";
          send_msg(wsi, msg, strlen(msg));
        }else{ /** there is no message being or to be processed */
          client_data->message_id = snowflake_id();

          /** append message id to the head of the raw message */
          size_t new_msg_len = len + MSG_ID_STR_SIZE;
          unsigned char * tmp_buf = malloc(new_msg_len);
          if (!tmp_buf) {
            log_to_file("Memory allocation failed for *message id appending in processing message form client a*");
            pthread_mutex_unlock(&clients_mutex);
            return -1;
          }
          memset(tmp_buf, 0, new_msg_len);
          sprintf(tmp_buf, "%020llu", client_data->message_id);
          memcpy(tmp_buf + MSG_ID_STR_SIZE, in, len);

          bool res = send_msg(b_client->wsi, tmp_buf, new_msg_len);

          free(tmp_buf);
          if (!res) {
            log_to_file("Failed to send message to B client");
            pthread_mutex_unlock(&clients_mutex);
            return -1;
          }

          client_data->msg_status = 2;
          waitting_for_gpt = true;
        }
      } else {
        /** tell the client the GPT Robot is offline */
        unsigned char * msg = "PesticideGPT Robot is offline!";
        send_msg(wsi, msg, strlen(msg));
      }
      pthread_mutex_unlock(&clients_mutex);
      break;
    case LWS_CALLBACK_CLOSED:
      pthread_mutex_lock(&clients_mutex);
      for (int i = 0; i < MAX_A_CLIENTS; ++i) {
        if (a_clients[i]->wsi == wsi) {
          if(NULL != a_clients[i]->msg_cache){
            free(a_clients[i]->msg_cache);
            a_clients[i]->msg_cache = NULL;
          }
          a_clients[i] = NULL;
          sprintf(log_message, "A client disconnected from slot %d", i);
          log_to_file(log_message);
          break;
        }
      }
      pthread_mutex_unlock(&clients_mutex);
      break;
    default:
      break;
  }
  return 0;
}

static int wss_b_callback(struct lws *wsi, enum lws_callback_reasons reason, void *user, void *in, size_t len) {
  client_data_t *client_data = (client_data_t *)user;
  char log_message[LOG_MSG_BUF_SIZE];
  memset(log_message, 0, sizeof(log_message));
  switch (reason) {
    case LWS_CALLBACK_ESTABLISHED:
      pthread_mutex_lock(&clients_mutex);
      if (b_client == NULL) {
        b_client = client_data;
        b_client->wsi = wsi;
        b_client->client_type = 1;
        sprintf(log_message, "B client connected");
      } else {
        /** disconnect */
        sprintf(log_message, "B client rejected (already connected)");
        unsigned char * msg = "Only one B client allowed";
        disconnect(wsi, log_message, msg);
        pthread_mutex_unlock(&clients_mutex);
        return -1;
      }
      log_to_file(log_message);
      pthread_mutex_unlock(&clients_mutex);
      break;
    case LWS_CALLBACK_RECEIVE:
      pthread_mutex_lock(&clients_mutex);
      uint64_t msg_id = peek_msg_id(in);
      for(int i=0; i<MAX_A_CLIENTS; ++i){
        if(a_clients[i]->wsi && a_clients[i]->message_id == msg_id){
          send_msg(a_clients[i]->wsi, in+MSG_ID_STR_SIZE, len-MSG_ID_STR_SIZE);
          a_clients[i]->message_id = 0L;
          a_clients[i]->msg_status = 0;
          if(a_clients[i]->msg_cache){
            free(a_clients[i]->msg_cache);
            a_clients[i]->msg_cache = NULL;
          }
          break;
        }
      }
      client_data_t * tmp = NULL;
      for(int i=0; i<MAX_A_CLIENTS; ++i){
        if(a_clients[i] && a_clients[i]->msg_status == 1){
          if(tmp){
            if(tmp->message_id>a_clients[i]->message_id){
              tmp = a_clients[i];
            }
          }else{
            tmp = a_clients[i];
          }
        }
      }
      if(tmp){
        /** append massage id to raw message */
        size_t new_msg_id = tmp->len + MSG_ID_STR_SIZE;
        unsigned char *tmp_buf = malloc(new_msg_id);
        if (!tmp_buf) {
          log_to_file("Memory allocation failed for *message id appending in processing message form client b*");
          pthread_mutex_unlock(&clients_mutex);
          return -1;
        }
        sprintf(tmp_buf, "%020llu", tmp->message_id);
        memcpy(tmp_buf + MSG_ID_STR_SIZE, tmp->msg_cache, tmp->len);

        /** send the message to B client */
        bool res = send_msg(b_client->wsi, tmp_buf, new_msg_id);

        free(tmp_buf);
        free(tmp->msg_cache);
        tmp->msg_cache = NULL;
        if (!res) {
          log_to_file("Failed to send message to B client");
        }

        tmp->msg_status = 2;
      }else{
        waitting_for_gpt = false;
      }
      pthread_mutex_unlock(&clients_mutex);
      break;
    case LWS_CALLBACK_CLOSED:
      pthread_mutex_lock(&clients_mutex);
      b_client = NULL;
      sprintf(log_message, "B client disconnected");
      log_to_file(log_message);
      pthread_mutex_unlock(&clients_mutex);
      break;
    default:
      break;
  }
  return 0;
}

static struct lws_protocols protocols[] = {
  {"wss-a", wss_a_callback, sizeof(client_data_t), LWS_RX_BUF_SIZE},
  {"wss-b", wss_b_callback, sizeof(client_data_t), LWS_RX_BUF_SIZE},
  {NULL, NULL, 0, 0}
};

void signal_handler(int signal_number) {
  terminate_server = 1;
  log_to_file("Termination signal received, shutting down...");
}

/** Cleaning work before safe exit */
void before_exit(){
  pthread_mutex_lock(&clients_mutex);
  /** Clean up Class A client connections and related resources */
  char log_message[LOG_MSG_BUF_SIZE];
  for(int i = 0; i < MAX_A_CLIENTS; ++i){
    if(a_clients[i] != NULL){
      if(a_clients[i]->msg_cache){
        free(a_clients[i]->msg_cache);
        a_clients[i]->msg_cache = NULL;
      }
      memset(log_message, 0, sizeof(log_message));
      sprintf(log_message, "Server actively exits");
      unsigned char * msg = "Server actively exits";
      if(a_clients[i]->wsi){
        disconnect(a_clients[i]->wsi, log_message, msg);
      }
    }
  }
  /** Clean up Class B client connections and related resources */
  if(b_client != NULL){
    memset(log_message, 0, sizeof(log_message));
    sprintf(log_message, "Server actively exits");
    unsigned char * msg = "Server actively exits";
    if(b_client->wsi){
      disconnect(b_client->wsi, log_message, msg);
    }
  }
  pthread_mutex_unlock(&clients_mutex);
}
/*
int read_config_file(const char *filename, int *port, char *ssl_cert, char *ssl_key) {
  FILE *file = fopen(filename, "r");
  if (!file) {
    return -1;
  }
  char section[64];
  memset(section, 0, sizeof(section));
  while (fscanf(file, "%63s", section) == 1) {
    if (strcmp(section, "[port]") == 0) {
      fscanf(file, "%d", port);
    } else if (strcmp(section, "[ssl_cert]") == 0) {
      fscanf(file, "%255s", ssl_cert);
    } else if (strcmp(section, "[ssl_key]") == 0) {
      fscanf(file, "%255s", ssl_key);
    }
    memset(section, 0, sizeof(section));
  }
  fclose(file);
  return 0;
}
*/
int main() {
  lws_set_log_level(LLL_ERR | LLL_WARN, NULL);

  struct lws_context_creation_info info;
  memset(&info, 0, sizeof(info));

  int port=8443;
  /*
  char ssl_cert_path[256];
  char ssl_key_path[256];
  if (read_config_file(CONFIG_FILE, &port, ssl_cert_path, ssl_key_path) != 0) {
    log_to_file("read confure file failed");
    return EXIT_FAILURE;
  }
  */
  info.port = port;
  info.protocols = protocols;
  info.gid = -1;
  info.uid = -1;
  /*
  info.options = LWS_SERVER_OPTION_DO_SSL_GLOBAL_INIT;
  info.ssl_cert_filepath = ssl_cert_path;
  info.ssl_private_key_filepath = ssl_key_path;
  */

  struct lws_context *context = lws_create_context(&info);
  if(!context){
    log_to_file("Failed to create libwebsockets context");
    return EXIT_FAILURE;
  }
  log_to_file("WebSocket server started on port 8443");

  signal(SIGUSR1, signal_handler);

  while(!terminate_server) {
    lws_service(context, LWS_TIMEOUT_MS);
  }

  before_exit();

  lws_context_destroy(context);
  log_to_file("WebSocket server stopped");

  return EXIT_SUCCESS;
}
